% Calculates the stationary distribution for the menu cost model
%
% Q_old is the initial probability distribution.
% 
% Jon Steinsson and Emi Nakamura, August 2006
%***************************************************************

function Q_new = StaDistGECalvoPlus(Q_old,alphaCalvo,Proba,ProbNgr,...
    F,F2,G,a,rp,rad,tol,yesprint)

agridsize = size(a,1);
rpgridsize = size(rp,1);
radgridsize = size(rad,1);
Ssize = agridsize*rpgridsize*radgridsize;

Q_old = reshape(Q_old,[Ssize 1]);
F = reshape(F,[Ssize 1]);
F2 = reshape(F2,[Ssize 1]);

% Calculate GInd
%**********************************************************************
% GInd is a vector that describes how dividing by inflation moves firms 
% in certain states to new states.

GInd = CalculateGInd(a,rp,rad,G);

% Calculate Find and FpcInd
%**********************************************************************
% FInd is a vector that describes how applying F moves firms in certain
% states to new states.

QaInd = repmat((1:agridsize)',rpgridsize*radgridsize,1);

QradInd = repmat((1:radgridsize)',[1 agridsize rpgridsize]);
QradInd = permute(QradInd,[2 3 1]);
QradInd = reshape(QradInd,[Ssize 1]);

[dummy,QrpInd] = ismember(int32(F*10^6),int32(rp*10^6));
clear dummy
FInd = sub2ind([agridsize rpgridsize radgridsize],QaInd,QrpInd,QradInd);

[dummy,QrpInd] = ismember(int32(F2*10^6),int32(rp*10^6));
clear dummy
F2Ind = sub2ind([agridsize rpgridsize radgridsize],QaInd,QrpInd,QradInd);
clear QaInd QrpInd QradInd

% Sort GInd, FInd and FpcInd
%**********************************************************************

OldF2Ind = (1:Ssize)';

[GInd iGInd] = sort(GInd);
OldGInd = OldF2Ind(iGInd);

[FInd iFInd] = sort(FInd);
OldFInd = OldF2Ind(iFInd);

[F2Ind iF2Ind] = sort(F2Ind);
OldF2Ind = OldF2Ind(iF2Ind);

% Calculate stationary distribution
%**********************************************************************

time23 = cputime;
Q_new = Q_old;
i = 1;
supDifQ = 2*tol;
while supDifQ > tol
    %fprintf('supDifQ = %5.10f  tol = %5.10f\n',supDifQ,tol)
    Q_old = Q_new;

    % Apply F
    Q_new = accumarray(FInd,Q_old(OldFInd));
    Q_new = [Q_new; zeros(size(Q_old,1)-size(Q_new,1),1)];
    Q2_new = accumarray(F2Ind,Q_old(OldF2Ind));
    Q2_new = [Q2_new; zeros(size(Q_old,1)-size(Q2_new,1),1)];
    Q_new = alphaCalvo*Q_new + (1-alphaCalvo)*Q2_new;
    
    % Apply Proba
    Q_new = ((reshape(Q_new,agridsize,rpgridsize*radgridsize)')*Proba)';
    
    % Apply ProbNgr
    Q_new = reshape(Q_new,[agridsize*rpgridsize radgridsize])*ProbNgr;
    Q_new = reshape(Q_new,[Ssize 1]);    
    
    % Apply G
    Q_new = accumarray(GInd,Q_new(OldGInd));
    Q_new = [Q_new; zeros(size(Q_old,1)-size(Q_new,1),1)];
    
    supDifQ = max(abs(Q_old-Q_new));
    
    if yesprint == 1
        if mod(i,5) == 0
            fprintf('$')
        end
        if mod(i,400) == 0
            fprintf('\n')
        end
    end
    i = i + 1;
end
if yesprint == 1
    fprintf('\n\n')
    fprintf('Time of Loop: %8.3f\n\n',cputime - time23)
end

    
    

